# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/25 16:34
# @Author : zxl
# @Email : zhangxl_njit@163.com
# @File : reinit_stream.py
# @Version: open-planner 2.0
"""
文件说明：重新定义流

"""
import math
import random
import json
import importlib
from OpenPlanner.core.src.op_class import OPClass
from OpenPlanner.core.api.algo_input_class import StreamInfo
from copy import *


def stream_reinit_hook(op_obj: OPClass):
    _init_link_stream(op_obj)
    stream_obj_set = op_obj.algo_input.stream_array
    link_obj_set = op_obj.algo_input.link_array
    algo_param_dict = op_obj.algo_param

    # _simple_random_win(stream_obj_set, link_obj_set, algo_param_dict)
    # 几种不同的方式生成起止窗口
    _trans_streams(stream_obj_set, link_obj_set, algo_param_dict)
    # _no_limit_win(stream_obj_set, link_obj_set)
    # 重写流文件
    _rewrite_input(op_obj)


# 起始，终止完全随机
def _trans_streams(stream_array, link_array, algo_param_dict):
    huge_slot = int(algo_param_dict['huge_slot'])
    slot_num = int(algo_param_dict['slot_num'])
    buf_num = int(algo_param_dict['buf_num'])
    for stream in stream_array:
        period = stream.period
        period_slot = math.ceil(period / huge_slot)
        latence = stream.latency_requirement
        # latence转换成huge_slot
        latence_slot = math.ceil(latence / huge_slot)
        route = stream.route
        # 该值可调，talker的右界，单位为时隙
        # slot_num*huge_slot<=latency
        # tx_win = [0, slot_num]
        # 起始窗口在1,slotnum中随机，终止窗口在符合约束下左右偏移shift_slot_num个窗口
        shift_slot_num = random.randint(0, slot_num)
        # start_slot = random.randint(0, math.ceil(shift_slot_num / 2))
        # end_slot = random.randint(start_slot, shift_slot_num)
        # shift_slot_num = slot_num
        # tx_win = [0, shift_slot_num]
        start_slot = random.randint(0, period_slot)
        end_slot = random.randint(start_slot, period_slot)
        tx_win = [start_slot, end_slot]
        sw_num = len(route) + 1
        # 根据最大最小延迟计算接收端可能的窗口范围
        # 然后向外扩展2*slot_num个窗口（左右各slot_num个）
        ND_min = sw_num + 1
        ND_max = sw_num * buf_num + buf_num
        # 接收窗口应满足链路的最长、最短传输约束
        rx_win_scope = [0] * 2
        rx_win_scope[0] = tx_win[0] + ND_min
        rx_win_scope[1] = tx_win[1] + ND_max
        rx_win = [0] * 2
        # 初始条件的接收窗口应该包含实际根据链路约束算出的窗口
        # 在这里将约束得到的窗口向两边扩充slot_num个时隙
        # listener的左窗口应大于talker的右窗口
        rx_win[0] = max(tx_win[1], rx_win_scope[0] - shift_slot_num)
        rx_win[1] = rx_win_scope[1] + shift_slot_num
        # if tx_win[1] + ND_max < rx_win[0] or tx_win[0] + ND_min > rx_win[1]:
        #     print("talker和listener窗口约束不合理！！")
        #     raise RuntimeError('窗口约束不合理，导致该算法无解')
        # 另外，listener的右窗口应符合最大延迟约束，且不能超过period
        listener_max = min(rx_win[1], latence_slot - tx_win[0], period_slot)
        rx_win[1] = listener_max
        talker_win = [0, tx_win[1] * huge_slot]
        listener_win = [rx_win[0] * huge_slot, rx_win[1] * huge_slot]
        talker_id = link_array[route[0]].src_node
        listener_id = link_array[route[-1]].dst_node
        opts = {'talker': {'node_id': talker_id, 'win': talker_win}, 'listeners': []}
        opts['listeners'].append({'listener': {'node_id': listener_id, 'win': listener_win}})
        stream.opts = opts
    # _rewrite_input(stream_array)


#  简单随机talker窗口和listener窗口
def _simple_random_win(stream_array, link_array, algo_param_dict):
    huge_slot = int(algo_param_dict['huge_slot'])
    slot_num = int(algo_param_dict['slot_num'])
    buf_num = int(algo_param_dict['buf_num'])
    for stream in stream_array:
        period = stream.period
        period_slot = math.ceil(period / huge_slot)
        latence = stream.latency_requirement
        # latence转换成huge_slot
        latence_slot = math.ceil(latence / huge_slot)
        route = stream.route
        # 该值可调，talker的右界，单位为时隙
        # slot_num*huge_slot<=latency
        # tx_win = [0, slot_num]
        # 起始窗口在1,slotnum中随机，终止窗口在符合约束下左右偏移shift_slot_num个窗口
        shift_slot_num = random.randint(1, slot_num)
        start_slot = random.randint(0, math.ceil(shift_slot_num / 2))
        end_slot = random.randint(start_slot, shift_slot_num)
        # shift_slot_num = slot_num
        # tx_win = [0, shift_slot_num]
        tx_win = [start_slot, end_slot]
        sw_num = len(route) + 1
        # 根据最大最小延迟计算接收端可能的窗口范围
        # 然后向外扩展2*slot_num个窗口（左右各slot_num个）
        ND_min = sw_num + 1
        ND_max = sw_num * buf_num + buf_num
        # 接收窗口应满足链路的最长、最短传输约束
        rx_win_scope = [0] * 2
        rx_win_scope[0] = tx_win[0] + ND_min
        rx_win_scope[1] = tx_win[1] + ND_max
        rx_win = [0] * 2
        # 初始条件的接收窗口应该包含实际根据链路约束算出的窗口
        # 在这里将约束得到的窗口向两边扩充slot_num个时隙
        # listener的左窗口应大于talker的右窗口
        rx_win[0] = max(tx_win[1], rx_win_scope[0] - shift_slot_num)
        rx_win[1] = rx_win_scope[1] + shift_slot_num
        # if tx_win[1] + ND_max < rx_win[0] or tx_win[0] + ND_min > rx_win[1]:
        #     print("talker和listener窗口约束不合理！！")
        #     raise RuntimeError('窗口约束不合理，导致该算法无解')
        # 另外，listener的右窗口应符合最大延迟约束，且不能超过period
        listener_max = min(rx_win[1], latence_slot - tx_win[0], period_slot)
        rx_win[1] = listener_max
        talker_win = [0, tx_win[1] * huge_slot]
        listener_win = [rx_win[0] * huge_slot, rx_win[1] * huge_slot]
        talker_id = link_array[route[0]].src_node
        listener_id = link_array[route[-1]].dst_node
        opts = {'talker': {'node_id': talker_id, 'win': talker_win}, 'listeners': []}
        opts['listeners'].append({'listener': {'node_id': listener_id, 'win': listener_win}})
        stream.opts = opts
    #     不放在这边写，只有有解的时候才写输入文件
    # _rewrite_input(stream_array)


def _no_limit_win(stream_array, link_array):
    for stream in stream_array:
        period = stream.period
        # latence转换成huge_slot
        route = stream.route
        talker_win = [0, period]
        listener_win = [0, period]
        talker_id = link_array[route[0]].src_node
        listener_id = link_array[route[-1]].dst_node
        opts = {'talker': {'node_id': talker_id, 'win': talker_win}, 'listeners': []}
        opts['listeners'].append({'listener': {'node_id': listener_id, 'win': listener_win}})
        stream.opts = opts


def _init_link_array(op_obj):
    # 从topo_file中读取link_set
    topo_filepath = op_obj.topo_filepath
    with open(topo_filepath, 'r') as fd:
        link_set = json.load(fd)
    # 判断用户是否是自定义Link类
    user_defined_link = op_obj.init_input.user_defined_link
    link_name = 'Link'
    LinkClass = getattr(
        importlib.import_module('OpenPlanner.core.api.algo_input_class'),
        link_name
    )
    if user_defined_link != '':
        # 动态加载自定义链路
        class_name = user_defined_link
        base_module = op_obj.init_input.base_module_path
        try:
            module = importlib.import_module(base_module)
        except ModuleNotFoundError:
            print('Error: 用户自定义模块%s不存在，'
                  '请检查路径后更新配置文件' % base_module)
            exit(0)
        try:
            LinkClass = getattr(module, class_name)
        except AttributeError:
            print('AttributeError: 用户自定义模块%s'
                  '没有%s属性，'
                  '请检查路径后更新配置文件'
                  % (base_module, class_name))
            exit(0)

    link_array = []
    for link in link_set:
        link_obj = LinkClass(**link)
        link_array.append(link_obj)

    op_obj.algo_input.link_array = link_array


def _init_stream_array(op_obj):
    stream_filepath = op_obj.stream_filepath
    with open(stream_filepath, 'r') as fd:
        stream_set = json.load(fd)
    # 判断用户是否自定义Stream类
    user_defined_stream = \
        op_obj.init_input.user_defined_stream
    StreamClass = getattr(
        importlib.import_module('OpenPlanner.core.api.algo_input_class'),
        'Stream'
    )
    if user_defined_stream != '':
        # 动态加载自定义流量类
        class_name = op_obj.init_input.user_defined_stream
        base_module = op_obj.init_input.base_module_path
        try:
            module = importlib.import_module(base_module)
        except ModuleNotFoundError:
            print('ModuleNotFoundError: 用户自定义流量模块'
                  '%s不存在，请检查模块路径后'
                  '更新配置文件' % base_module)
            exit(0)

        try:
            StreamClass = getattr(module, class_name)
        except AttributeError:
            print('AttributeError: 用户自定义模块%s'
                  '中不存在%s类，请检查路径后更新配置文件'
                  % (base_module, class_name))
            exit(0)

    stream_array = []
    for stream in stream_set:
        stream_obj = StreamClass(**stream)
        stream_array.append(stream_obj)
        hop_id = 0
        for link_id in stream['route']:
            link_obj = \
                op_obj.algo_input.link_array[link_id]
            link_obj.add_stream_to_link(
                StreamInfo(
                    **{"stream_id": stream['stream_id'],
                       "hop_id": hop_id}
                )
            )
            hop_id += 1
    op_obj.algo_input.stream_array = stream_array


def _init_link_stream(op_obj: OPClass):
    _init_link_array(op_obj)
    _init_stream_array(op_obj)


# 重写stream的输入文件
def _rewrite_input(op_obj):
    input_stream_file = op_obj.stream_filepath
    stream_array = op_obj.algo_input.stream_array
    f = open(input_stream_file, 'w')
    streams_dict = []
    stream_obj_set = deepcopy(stream_array)
    for stream in stream_obj_set:
        stream_dict = stream.__dict__
        # 重新拼接输出
        streams_dict.append(stream_dict)
    streams_json = json.dumps(streams_dict, indent=4, ensure_ascii=False)
    f.write(streams_json)
    f.close()
